from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import datetime  # For datetime objects
import os.path  # To manage paths
import sys  # To find out the script name (in argv[0])
sys.path.append('./backtrader')

# Import the backtrader platform
import backtrader as bt

#手动执行,会打印交易日志 ,设置为false ,可以自动回测所有组合参数,找到年化最大的参数组合,填写到param后,手动执行
#手动执行是会打印所有的交易日志
manu_test_en = True
# manu_test_en = False

# Create a Stratey
class TestStrategy(bt.Strategy):
    params = (
        ('maperiod_2', 2),
        ('maperiod_5', 5),
        ('maperiod_10', 10),
        ('maperiod_20', 20),
        ('maperiod_30', 30),
        ('divergence_period', 80),  # 底背离判断的天数
        ('top_divergence_period', 60),  # 顶背离判断的天数

        ('data_path_global', './csv_data/SH_000300_stock_data.csv'), 
        ('bollinger_period', 55),  # 布林带周期 000300 沪深300
        ('bollinger_dev', 2.0),    # 布林带标准差倍数 

        # ('data_path_global', './csv_data/sh_601818_stock_data.csv'), 
        # ('bollinger_period', 10),  # 布林带周期  601818 光大银行
        # ('bollinger_dev', 1.0),    # 布林带标准差倍数

        # ('data_path_global', './csv_data/sz_300498_stock_data.csv'), 
        # ('bollinger_period', 45),  # 布林带周期 300498  温氏股份
        # ('bollinger_dev', 2.0),    # 布林带标准差倍数 


        ('buy_amount', 20000),      # 每次买入固定金额
        ('max_buy_times', 5),       # 最大买入次数
    )

    def log(self, txt, dt=None):
        ''' Logging function fot this strategy'''
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    def __init__(self):
        # Keep a reference to the "close" line in the data[0] dataseries
        self.dataclose = self.datas[0].close
        # 引用成交量数据
        self.volume = self.datas[0].volume
       
        # To keep track of pending orders and buy price/commission
        self.order = None
        self.buyprice = None
        self.buycomm = None
        
        # 统计变量
        self.trade_count = 0  # 交易次数
        self.start_cash = None  # 初始资金
        self.start_date = None  # 开始日期
        self.end_date = None  # 结束日期
        self.initial_cash = 100000.0  # 固定的起始资金基数
        self.total_cost = 0.0
        # 买入记录相关变量
        self.buy_count = 0  # 买入次数计数器
        self.buy_shares_list = []  # 记录每次买入的股数

        # Add a MovingAverageSimple indicator
        self.sma = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.maperiod_5)
        self.sma10 = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.maperiod_10)
        self.sma20 = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.maperiod_20)

        # 添加MACD指标
        # self.macd_histo = bt.indicators.MACDHisto(self.datas[0])
        # add diff and dea indicator
        self.macd = bt.indicators.MACD(self.datas[0])
        # MACD包含了三个值：macd (DIFF线), signal (DEA线), histo (柱状图)
        self.diff = self.macd.macd  # DIFF线 (快速EMA - 慢速EMA)
        self.dea = self.macd.signal  # DEA线 (DIFF的EMA平滑线)
        # 添加最低价指标用于底背离判断
        self.lowest_period = bt.indicators.Lowest(self.datas[0].low, period=self.params.divergence_period)
        
        # 添加最高价指标用于顶背离判断
        self.highest_period = bt.indicators.Highest(self.datas[0].high, period=self.params.top_divergence_period)
        
        # 添加布林带指标
        self.bollinger = bt.indicators.BollingerBands(
            self.datas[0],
            period=self.params.bollinger_period,
            devfactor=self.params.bollinger_dev
        )

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            # Buy/Sell order submitted/accepted to/by broker - Nothing to do
            return

        # Check if an order has been completed
        # Attention: broker could reject order if not enough cash
        if order.status in [order.Completed]:
            if order.isbuy():
                # 记录买入的股数
                buy_shares = order.executed.size
                self.buy_shares_list.append(buy_shares)
                self.buy_count += 1
                
                # 累计总成本（含手续费）
                self.total_cost += order.executed.value + order.executed.comm
                
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
                
                # 买入后输出当前持仓资金信息
                current_value = self.broker.getvalue()
                current_cash = self.broker.getcash()
                position_size = self.position.size
                position_value = position_size * order.executed.price
                if manu_test_en:
                    self.log('BUY EXECUTED - Price: %.2f, 买入股数: %d, 成本: %.2f, 手续费: %.2f, 总资产: %.2f, 现金: %.2f, 持仓数量: %d, 持仓市值: %.2f, 买入次数: %d' %
                             (order.executed.price, buy_shares, order.executed.value, order.executed.comm, current_value, current_cash, position_size, position_value, self.buy_count))
                
            else:  # Sell
                # 按FIFO原则卖出：移除已卖出的股数记录
                sell_shares = abs(order.executed.size)  # 使用绝对值获取实际卖出股数
                remaining_sell = sell_shares
                
                # 从买入记录列表中依次扣除卖出的股数
                while remaining_sell > 0 and self.buy_shares_list:
                    if self.buy_shares_list[0] <= remaining_sell:
                        # 第一笔买入记录的股数完全被卖出
                        remaining_sell -= self.buy_shares_list[0]
                        self.buy_shares_list.pop(0)  # 移除这笔买入记录
                    else:
                        # 第一笔买入记录部分被卖出
                        self.buy_shares_list[0] -= remaining_sell
                        remaining_sell = 0
                
                # 卖出后输出当前持仓资金信息
                current_value = self.broker.getvalue()
                current_cash = self.broker.getcash()
                position_size = self.position.size
                position_value = position_size * order.executed.price if position_size > 0 else 0
                if manu_test_en:
                    self.log('SELL EXECUTED - Price: %.2f, 卖出股数: %d, 成本: %.2f, 手续费: %.2f, 总资产: %.2f, 现金: %.2f, 持仓数量: %d, 持仓市值: %.2f' %
                             (order.executed.price, sell_shares, abs(order.executed.value), order.executed.comm, current_value, current_cash, position_size, position_value))

            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            # self.log('Order Canceled/Margin/Rejected=%s' % order.status)
            pass

        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        # 统计交易次数
        self.trade_count += 1
        
        # 如果没有持仓了，重置买入计数器
        if self.position.size == 0:
            self.buy_count = 0
            self.buy_shares_list = []
        
        # self.log('OPERATION PROFIT, GROSS %.2f, NET %.2f' %
        #          (trade.pnl, trade.pnlcomm))


    def start(self):
        # 记录开始信息
        self.start_cash = self.broker.getvalue()
        # 暂时将start_date设为None，稍后在next方法中设置
        self.start_date = None
        self.log('策略开始运行 - 初始资金: %.2f' % self.start_cash)
        
    def stop(self):
        # 记录结束信息
        self.end_date = self.datas[0].datetime.date(0)
        final_cash = self.broker.getvalue()
        self.log('策略结束运行 - 结束日期: %s, 最终资金: %.2f' % (self.end_date, final_cash))
        
        # 计算统计数据
        total_return = ((final_cash - self.start_cash) / self.start_cash) * 100
        
        # 计算交易周期
        if self.start_date and self.end_date:
            trade_period = (self.end_date - self.start_date).days
            # 确保交易周期不为负数
            if trade_period < 0:
                trade_period = 0
        else:
            # 如果没有有效的开始和结束日期，使用默认值0
            trade_period = 0
        

        self.log('>>>>>>>>>>>>>>>finish:>>>>>>>>===' )
        if trade_period > 0:
            self.log('交易周期: %s 至 %s (共%d天), period=%d,bll_dev=%d,btm_diff=%d' % 
            (self.start_date, self.end_date, trade_period, self.params.bollinger_period, self.params.bollinger_dev, self.params.divergence_period))
            self.log('start_cash: %.2f, last_cash: %.2f, tatal_perc: %.2f%%, trade_ct: %d, yearprofit: %.2f%%' % 
            (self.start_cash, final_cash, total_return, self.trade_count, total_return * 365 / trade_period))
        else:
            self.log('===============trade none===============')

        self.log('================end==================')
   
        
    def check_buy_bollinger(self):
        """基于布林带的买入判断函数
        当价格触碰或跌破布林带下轨时返回1，否则返回0
        """
        # 确保有足够的数据计算布林带指标
        if len(self) >= self.params.bollinger_period:
            # 获取当前价格和布林带下轨值
            current_price = self.dataclose[0]
            lower_band = self.bollinger.lines.bot[0]
            
            # 判断价格是否触碰或跌破布林带下轨
            # 考虑一个小的容忍度，避免因价格微小波动而频繁触发信号
            if current_price <= lower_band * 1.005:
                return 1
        return 0
        
    def check_sell_bollinger(self):
        """基于布林带的卖出判断函数
        当价格触碰或突破布林带上轨时返回1，否则返回0
        """
        # 确保有足够的数据计算布林带指标
        if len(self) >= self.params.bollinger_period:
            # 获取当前价格和布林带上轨值
            current_price = self.dataclose[0]
            upper_band = self.bollinger.lines.top[0]
            
            # 判断价格是否触碰或突破布林带上轨
            # 考虑一个小的容忍度，避免因价格微小波动而频繁触发信号
            if current_price >= upper_band * 0.995:
                return 1
        return 0
        
    def check_buymacd(self):
        """基于MACD的买入判断函数
        判断DIFF和DEA在零轴下方形成金叉的条件
        """
        # 金叉：DIFF线从下方穿越DEA线
        # 零轴下方：当前DIFF和DEA都小于0
        # 需要至少２个数据点来判断穿越
        if (len(self) > 1 and 
            self.diff[0] < 0 and self.dea[0] < 0 and  # 当前都在零轴下方
            self.diff[0] > self.dea[0] and  # 当前DIFF在DEA上方
            self.diff[-1] <= self.dea[-1]):  # 前一根K线DIFF在DEA下方或相等
            return 1
        return 0
    
    def check_sell_Deviation(self):
        """顶背离卖出判断函数
        当价格达到指定期间最高价但DIFF没有创新高时（顶背离）返回1，否则返回0
        """
        # 需要至少指定天数的数据点来计算最高价
        if len(self) >= self.params.top_divergence_period:
            # 当前价格是否接近指定期间最高价（考虑小误差）
            current_high = self.datas[0].high[0]
            highest_period_value = self.highest_period[0]
            
            # 容忍小误差，判断当前高价是否接近指定期间最高价
            if abs(current_high - highest_period_value) < 0.01:
                # 计算指定期间内DIFF的最高值
                diff_period_max = float('-inf')
                for i in range(min(self.params.top_divergence_period, len(self))):
                    if self.diff[-i] > diff_period_max:
                        diff_period_max = self.diff[-i]
                
                # 判断当前DIFF是否低于指定期间内的最高值（顶背离）
                if self.diff[0] < diff_period_max:
                        return 1
        return 0
    def check_buy_Deviation(self):
        """底背离买入判断函数
        当价格达到指定期间最低价但DIFF没有创新低时（底背离）执行买入
        """
        # 需要至少指定天数的数据点来计算最低价
        if len(self) >= self.params.divergence_period:
            # 当前价格是否等于指定期间最低价
            current_low = self.datas[0].low[0]
            lowest_period_value = self.lowest_period[0]
            
            # 容忍小误差，判断当前低价是否接近指定期间最低价
            if abs(current_low - lowest_period_value) < 0.01:
                # 计算指定期间内DIFF的最低值
                diff_period_min = float('inf')
                for i in range(min(self.params.divergence_period, len(self))):
                    if self.diff[-i] < diff_period_min:
                        diff_period_min = self.diff[-i]
                
                # 判断当前DIFF是否高于指定期间内的最低值（底背离）
                if self.diff[0] > diff_period_min:
                    # self.log('*** 底背离信号：价格创%d日新低但DIFF未创新低！Price: %.2f, DIFF: %.2f, %d日DIFF最低: %.2f ***' % 
                            # (self.params.divergence_period, current_low, self.diff[0], self.params.divergence_period, diff_period_min))
                    # 如果没有仓位，执行买入操作
                    # if not self.position:
                    # self.order = self.buy()
                    return 1
        return 0

        
    def check_buymacd_diff(self):
        """基于DIFF线穿越零轴的买入判断函数
        当DIFF线上穿零轴时返回1，否则返回0
        """
        if (len(self) > 1 and 
            self.diff[0] > 0 and  # 当前DIFF在零轴上方
            self.diff[-1] <= 0):  # 前一根K线DIFF在零轴下方或等于零轴
            return 1
        return 0
        
    def check_macd_diff_shrink_below_zero(self):
        """判断MACD零轴下方DIFF三天收缩的条件
        当DIFF在零轴下方且连续三天数值逐渐收缩（绝对值减小）时返回1
        返回值：1-满足条件，0-不满足条件
        """
        # 需要至少三个数据点来判断三天的趋势
        if len(self) >= 3:
            # 获取近三天的DIFF值
            diff_today = self.diff[0]      # 今天
            diff_yesterday = self.diff[-1]  # 昨天
            diff_day_before = self.diff[-2] # 前天
            
            # 判断所有DIFF值都在零轴下方
            if diff_today < 0 and diff_yesterday < 0 and diff_day_before < 0:
                # 判断连续三天DIFF的绝对值逐渐减小（数值逐渐向零轴靠近）
                # abs(diff_today) < abs(diff_yesterday) < abs(diff_day_before)
                if (abs(diff_today) < abs(diff_yesterday) and 
                    abs(diff_yesterday) < abs(diff_day_before)):
                    return 1
        return 0
    def next(self):
        # 记录第一个交易日作为开始日期
        if self.start_date is None:
            self.start_date = self.datas[0].datetime.date(0)
            self.log('策略开始日期记录: %s' % self.start_date)
            
        # Simply log the closing price of the series from the reference
        # self.log('Close_day: %.2f, sma: %.2f, sma10: %.2f, sma20: %.2f, DIFF: %.2f, DEA: %.2f, macd_histo: %.2f' % 
        #         (self.dataclose[0], self.sma[0], self.sma10[0], self.sma20[0], self.diff[0], self.dea[0], self.macd_histo.histo[0]))
                # DIFF和DEA指标已添加并打印 
        
        
        
        # Check if an order is pending ... if yes, we cannot send a 2nd one
        if self.order:
            return
            
        #==================================================================
        # # MACD买入信号检查
        # if self.check_buymacd():  # 调用金叉买入判断函数
        #     # 检查是否达到最大买入次数
        #     if self.buy_count < self.params.max_buy_times:
        #         # 计算买入数量：根据固定金额和当前价格计算股数
        #         buy_amount = self.params.buy_amount
        #         current_price = self.dataclose[0]
        #         # 计算可买入的股数（整手）
        #         buy_shares = int(buy_amount / current_price / 100) * 100  # 整手买入（100股为一手）
                
        #         if buy_shares > 0:
        #             actual_amount = buy_shares * current_price
        #             if manu_test_en:
        #                 self.log('MACD金叉买入信号 - Price: %.5f, DIFF: %.5f, DEA: %.5f, 计划金额: %.2f, 实际金额: %.2f, 买入股数: %d, 第%d次买入' % 
        #                         (current_price, self.diff[0], self.dea[0], buy_amount, actual_amount, buy_shares, self.buy_count + 1))
        #             self.order = self.buy(size=buy_shares)
        #     else:
        #         if manu_test_en:
        #             self.log('已达到最大买入次数%d次，不再买入' % self.params.max_buy_times)
        
        # elif self.check_buymacd_diff():  # 调用DIFF线上穿零轴买入判断函数
        #     # 检查是否达到最大买入次数
        #     if self.buy_count < self.params.max_buy_times:
        #         # 计算买入数量：根据固定金额和当前价格计算股数
        #         buy_amount = self.params.buy_amount
        #         current_price = self.dataclose[0]
        #         # 计算可买入的股数（整手）
        #         buy_shares = int(buy_amount / current_price / 100) * 100  # 整手买入（100股为一手）
                
        #         if buy_shares > 0:
        #             actual_amount = buy_shares * current_price
        #             if manu_test_en:
        #                 self.log('MACD DIFF上穿零轴买入信号 - Price: %.5f, DIFF: %.5f, DEA: %.5f, 计划金额: %.2f, 实际金额: %.2f, 买入股数: %d, 第%d次买入' % 
        #                         (current_price, self.diff[0], self.dea[0], buy_amount, actual_amount, buy_shares, self.buy_count + 1))
        #             self.order = self.buy(size=buy_shares)
        #     else:
        #         if manu_test_en:
        #             self.log('已达到最大买入次数%d次，不再买入' % self.params.max_buy_times)
        # #==================================================================

        # 调用布林带买入判断函数
        # if self.check_buy_bollinger() and self.check_macd_diff_shrink_below_zero():
        if  self.check_buy_Deviation():
            # 检查是否达到最大买入次数
            if self.buy_count < self.params.max_buy_times:
                # 计算买入数量：根据固定金额和当前价格计算股数
                buy_amount = self.params.buy_amount
                current_price = self.dataclose[0]
                # 计算可买入的股数（整手）
                buy_shares = int(buy_amount / current_price / 100)*100 # 整手买入（100股为一手）
                
                if buy_shares > 0:
                    actual_amount = buy_shares * current_price
                    if manu_test_en:
                        self.log('boll买入信号 - Price: %.5f, Lower Band: %.5f, 计划金额: %.2f, 实际金额: %.2f, 买入股数: %d, 第%d次买入' % 
                                (current_price, self.bollinger.lines.bot[0], buy_amount, actual_amount, buy_shares, self.buy_count + 1))
                    self.order = self.buy(size=buy_shares)
            else:
                if manu_test_en:
                    self.log('已达到最大买入次数%d次，不再买入' % self.params.max_buy_times)
        
    

        # ====================卖出操作================
        if self.position and self.position.size > 0:
            # 优先检查止损信号
            # if self.check_stop_loss():
            #     # 止损逻辑：一次性全部清仓
            #     current_position = self.position.size
            #     current_market_value = current_position * self.dataclose[0]
            #     loss_pct = (self.total_cost - current_market_value) / self.total_cost * 100
                
            #     if manu_test_en:
            #         self.log('止损触发 - Price: %.5f, 亏损率: %.2f%%, 一次性全部清仓: %d股 *** STOP LOSS ***' % 
            #                 (self.dataclose[0], loss_pct, current_position))
            #     self.order = self.sell(size=current_position)  # 全部清仓
            
            # 布林带卖出信号（仅在没有止损时检查）
            if self.check_sell_bollinger():
                # 卖出逻辑：依次按买入数量卖出（FIFO）
                if self.buy_shares_list:
                    # 卖出最早买入的那批股票
                    sell_shares = self.buy_shares_list[0]
                    
                    if manu_test_en:
                        self.log('boll卖出信号 - Price: %.5f, Upper Band: %.5f, 依次卖出第1批买入股票: %d股 ***' % 
                                (self.dataclose[0], self.bollinger.lines.top[0], sell_shares))
                    self.order = self.sell(size=sell_shares)  # 卖出第一批买入的股数
            
            # MACD死叉卖出信号（仅在没有止损和布林带信号时检查）
            # elif self.check_sell_macd():
            #     # 卖出逻辑：依次按买入数量卖出（FIFO）
            #     if self.buy_shares_list:
            #         # 卖出最早买入的那批股票
            #         sell_shares = self.buy_shares_list[0]
                    
            #         if manu_test_en:
            #             self.log('MACD死叉卖出信号 - Price: %.5f, DIFF: %.5f, DEA: %.5f, 依次卖出第1批买入股票: %d股 ***' % 
            #                     (self.dataclose[0], self.diff[0], self.dea[0], sell_shares))
            #         self.order = self.sell(size=sell_shares)  # 卖出第一批买入的股数


if __name__ == '__main__':
    # Create a cerebro entity
    cerebro = bt.Cerebro()

    # Add a strategy
    if manu_test_en:
        cerebro.addstrategy(TestStrategy)
    else:
        # strats = cerebro.optstrategy(TestStrategy, bollinger_period=range(10, 90, 5), bollinger_dev=range(1, 4)) 
        strats = cerebro.optstrategy(TestStrategy, divergence_period=range(20, 120, 5)) 


    # Datas are in a subfolder of the samples. Need to find where the script is
    # because it could have been called from anywhere
    modpath = os.path.dirname(os.path.abspath(sys.argv[0]))
    datapath = os.path.join(modpath, TestStrategy.params.data_path_global)

    # Create a Data Feed
    data = bt.feeds.YahooFinanceCSVData(
        dataname=datapath,
        # Do not pass values before this date
        fromdate=datetime.datetime(2020, 8, 3),
        # fromdate=datetime.datetime(2021, 8, 3),
        # fromdate=datetime.datetime(2022, 8, 3),
        # fromdate=datetime.datetime(2023, 8, 3),
        # fromdate=datetime.datetime(2024, 8, 3),
        # Do not pass values after this date
        todate=datetime.datetime(2025, 9, 18),
        # todate=datetime.datetime(2025, 9, 18),
        # Do not pass values after this date
        reverse=False)

    # Add the Data Feed to Cerebro
    cerebro.adddata(data)

    # Set our desired cash start
    cerebro.broker.setcash(100000.0)

    # 不使用固定仓位大小，改用策略内部的百分比仓位管理
    # cerebro.addsizer(bt.sizers.FixedSize, stake=5)

    # Set the commission - 万5手续费 (0.05%)
    cerebro.broker.setcommission(commission=0.0005)

    # Print out the starting conditions
    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())

    # Run over everything
    cerebro.run()

    # Print out the final result
    final_value = cerebro.broker.getvalue()
    print('Final Portfolio Value: %.2f' % final_value)
    
    
    # Plot the result
    if manu_test_en:
        cerebro.plot()